
import java.io.*;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Date;
import java.util.List;

public class AES {

	public static void main(String[] args) throws Exception {
		AES.encryptFile("c:\\data\\FileEncryptionTest.java");

		if (args.length > 0) {
			List<String> filesList = new ArrayList<String>();
			File file = new File(args[0]);
			File[] tempList = file.listFiles();
			for (int i = 0; i < tempList.length; i++) {
				if (tempList[i].isFile()) {
					filesList.add(tempList[i].toString());
				}
			}
			for (int i = 0; i < filesList.size(); i++) {
				AES.encryptFile(filesList.get(i));
			}
		}
	}

	byte[] word;
	// 轮数 当Nb=4时对应Nk取值4、6、8，Nr取值分别为10、12、14
	final static int Nr = 10;
	// 数据块字数 定值
	final static int Nb = 4;
	// 密钥字数Nk ES算法中规定密钥长度为128位、192位或256位，即Nk取值为4、6或8。
	final static int Nk = 4;

	private static final String KEY = "00012001710198aeda79171460153594";

	public AES(byte[] b) {
		word = new byte[4];
		for (int i = 0; i < 4; i++)
			word[i] = b[i];
	}

	public AES(AES w) {
		word = new byte[4];
		for (int i = 0; i < 4; i++)
			word[i] = w.word[i];
	}

	@Override
	public String toString() {
		String str = "";
		for (byte b : word)
			str += Integer.toHexString((b & 0xff) + 0x100).substring(1);
		return str;
	}

	/**
	 * 在GF(2^8)上的多项式加法
	 *
	 * @param a
	 * @param b
	 * @return
	 */
	static AES add(AES a, AES b) {
		AES c = new AES(new byte[4]);
		for (int i = 0; i < 4; i++)
			c.word[i] = add(a.word[i], b.word[i]);
		return c;
	}

	/**
	 * 在GF(2^8)上的多项式乘法
	 *
	 * @param a
	 * @param b
	 * @return
	 */
	static AES multiply(AES a, AES b) {
		AES c = new AES(new byte[4]);
		c.word[0] = add(add(add(multiply(a.word[0], b.word[0]), multiply(a.word[3], b.word[1])),
				multiply(a.word[2], b.word[2])), multiply(a.word[1], b.word[3]));
		c.word[1] = add(add(add(multiply(a.word[1], b.word[0]), multiply(a.word[0], b.word[1])),
				multiply(a.word[3], b.word[2])), multiply(a.word[2], b.word[3]));
		c.word[2] = add(add(add(multiply(a.word[2], b.word[0]), multiply(a.word[1], b.word[1])),
				multiply(a.word[0], b.word[2])), multiply(a.word[3], b.word[3]));
		c.word[3] = add(add(add(multiply(a.word[3], b.word[0]), multiply(a.word[2], b.word[1])),
				multiply(a.word[1], b.word[2])), multiply(a.word[0], b.word[3]));
		return c;
	}

	/**
	 * 在GF(2^8)上的多项式倍乘
	 *
	 * @param a
	 * @return
	 */
	static AES xtime(AES a) {
		AES b = new AES(new byte[4]);
		for (int i = 0; i < 4; i++)
			b.word[i] = a.word[(i + 1) % 4];
		return b;
	}

	static int m = 0x11b; // m=100011011

	/**
	 * 在GF(2^8)上的加法
	 * 
	 * @param a
	 * @param b
	 * @return
	 */
	static byte add(byte a, byte b) {
		return (byte) (a ^ b);
	}

	/**
	 * 在GF(2^8)上的求模
	 * 
	 * @param a
	 * @param b
	 * @return
	 */
	static byte mod(int a, int b) {
		String str_a = Integer.toBinaryString(a);
		String str_b = Integer.toBinaryString(b);
		if (str_a.length() < str_b.length())
			return (byte) a;
		return mod(a ^ (b << (str_a.length() - str_b.length())), b);
	}

	/**
	 * 在GF(2^8)上的乘法
	 * 
	 * @param a
	 * @param b
	 * @return
	 */
	static byte multiply(byte a, byte b) {
		int op = a & 0xff;
		char[] c = Integer.toBinaryString((b & 0xff) + 0x100).substring(1).toCharArray();
		int r = 0;
		for (int i = 0; i < c.length; i++)
			if (c[i] == '1')
				r ^= op << (7 - i);
		return mod(r, m);
	}

	/**
	 * 在GF(2^8)上的乘法逆
	 * 
	 * @param a
	 * @return
	 */
	static byte inverse(byte a) {
		if (a == 0)
			return 0;
		byte b = -128;
		while (mod(multiply(a, b), m) != 1)
			b++;
		return b;
	}

	/**
	 * 在GF(2^8)上的倍乘
	 * 
	 * @param a
	 * @return
	 */
	static byte xtime(byte a) {
		int r = (a & 0xff) << 1;
		if (r > 127)
			return mod(r, m);
		return (byte) r;
	}

	/**
	 * 仿射变换
	 * 
	 * @param b
	 * @param sign C：加密 D：解密
	 * @return
	 */
	static byte AffineTransformation(byte b, char sign) {
		byte[] x = Integer.toBinaryString((b & 0xff) + 0x100).substring(1).getBytes();
		for (int i = 0; i < x.length; i++)
			x[i] -= '0';
		if (sign == 'C') {
			byte[] x_ = new byte[8];
			byte b_ = 0;
			for (int i = 0; i < 8; i++) {
				x_[i] = (byte) (x[i] ^ x[(i + 1) % 8] ^ x[(i + 2) % 8] ^ x[(i + 3) % 8] ^ x[(i + 4) % 8]);
				b_ += x_[i] * Math.pow(2, 7 - i);
			}
			return (byte) (b_ ^ 0x63);
		} else {
			byte[] x_ = new byte[8];
			byte b_ = 0;
			for (int i = 0; i < 8; i++) {
				x_[i] = (byte) (x[(i + 1) % 8] ^ x[(i + 3) % 8] ^ x[(i + 6) % 8]);
				b_ += x_[i] * Math.pow(2, 7 - i);
			}
			return (byte) (b_ ^ 0x05);
		}
	}

	/**
	 * 加密密钥扩展
	 * 
	 * @param CipherKey
	 * @return
	 */
	static AES[][] KeyExpansion(AES[] CipherKey) {
		AES[] W = new AES[Nb * (Nr + 1)];
		// 密钥扩展
		AES Temp;
		if (Nk <= 6) {
			for (int i = 0; i < Nk; i++)
				W[i] = CipherKey[i];
			for (int i = Nk; i < W.length; i++) {
				Temp = new AES(W[i - 1]);
				if (i % Nk == 0)
					Temp = add(SubByte(Rotl(Temp)), Rcon(i / Nk));
				W[i] = add(W[i - Nk], Temp);
			}
		} else {
			for (int i = 0; i < Nk; i++)
				W[i] = CipherKey[i];
			for (int i = Nk; i < W.length; i++) {
				Temp = new AES(W[i - 1]);
				if (i % Nk == 0)
					Temp = add(SubByte(Rotl(Temp)), Rcon(i / Nk));
				else if (i % Nk == 4)
					Temp = SubByte(Temp);
				W[i] = add(W[i - Nk], Temp);
			}
		}
		// 轮密钥选择
		AES[][] RoundKey = new AES[Nr + 1][Nb];
		for (int i = 0; i < Nr + 1; i++)
			for (int j = 0; j < Nb; j++)
				RoundKey[i][j] = W[Nb * i + j];
		return RoundKey;
	}

	static AES SubByte(AES a) {
		AES w = new AES(a);
		for (int i = 0; i < 4; i++) {
			// 乘法逆代替
			w.word[i] = inverse(w.word[i]);
			// 仿射变换
			w.word[i] = AffineTransformation(w.word[i], 'C');
		}
		return w;
	}

	static AES Rotl(AES a) {
		AES w = new AES(a);
		byte b = w.word[0];
		for (int i = 0; i < 3; i++)
			w.word[i] = w.word[i + 1];
		w.word[3] = b;
		return w;
	}

	static AES Rcon(int n) {
		AES Rcon = new AES(new byte[4]);
		byte RC = 1;
		for (int i = 1; i < n; i++)
			RC = xtime(RC);
		Rcon.word[0] = RC;
		return Rcon;
	}

	/**
	 * 解密密钥扩展
	 *
	 * @param CipherKey
	 * @return
	 */
	static AES[][] InvKeyExpansion(AES[] CipherKey) {
		AES[][] InvRoundKey = KeyExpansion(CipherKey);
		for (int i = 1; i < Nr; i++)
			InvRoundKey[i] = InvMixColumn(InvRoundKey[i]);
		return InvRoundKey;
	}

	/**
	 * S盒变换
	 *
	 * @param state
	 * @return
	 */
	static AES[] ByteSub(AES[] state) {
		for (int i = 0; i < Nb; i++)
			for (int j = 0; j < 4; j++) {
				// 乘法逆代替
				state[i].word[j] = inverse(state[i].word[j]);
				// 仿射变换
				state[i].word[j] = AffineTransformation(state[i].word[j], 'C');
			}
		return state;
	}

	/**
	 * S盒逆变换
	 *
	 * @param state
	 * @return
	 */
	static AES[] InvByteSub(AES[] state) {
		for (int i = 0; i < Nb; i++)
			for (int j = 0; j < 4; j++) {
				// 仿射变换
				state[i].word[j] = AffineTransformation(state[i].word[j], 'D');
				// 乘法逆代替
				state[i].word[j] = inverse(state[i].word[j]);
			}
		return state;
	}

	/**
	 * 行移位变换
	 *
	 * @param state
	 * @return
	 */
	static AES[] ShiftRow(AES[] state) {
		byte[][] b = new byte[4][Nb];
		for (int j = 0; j < Nb; j++)
			for (int i = 0; i < 4; i++)
				b[i][j] = state[j].word[i];
		for (int i = 1; i < 4; i++)
			for (int k = 0; k < i; k++) {
				byte t = b[i][0];
				for (int j = 0; j < Nb - 1; j++)
					b[i][j] = b[i][j + 1];
				b[i][Nb - 1] = t;
			}
		for (int j = 0; j < Nb; j++)
			for (int i = 0; i < 4; i++)
				state[j].word[i] = b[i][j];
		return state;
	}

	/**
	 * 行移位逆变换
	 *
	 * @param state
	 * @return
	 */
	static AES[] InvShiftRow(AES[] state) {
		byte[][] b = new byte[4][Nb];
		for (int j = 0; j < Nb; j++)
			for (int i = 0; i < 4; i++)
				b[i][j] = state[j].word[i];
		for (int i = 1; i < 4; i++)
			for (int k = 0; k < Nb - i; k++) {
				byte t = b[i][0];
				for (int j = 0; j < Nb - 1; j++)
					b[i][j] = b[i][j + 1];
				b[i][Nb - 1] = t;
			}
		for (int j = 0; j < Nb; j++)
			for (int i = 0; i < 4; i++)
				state[j].word[i] = b[i][j];
		return state;
	}

	/**
	 * 列混合变换
	 *
	 * @param state
	 * @return
	 */
	static AES[] MixColumn(AES[] state) {
		byte[] b = { (byte) 0x02, (byte) 0x01, (byte) 0x01, (byte) 0x03 };
		AES a = new AES(b);
		for (int i = 0; i < Nb; i++)
			state[i] = multiply(a, state[i]);
		return state;
	}

	/**
	 * 列混合逆变换
	 *
	 * @param state
	 * @return
	 */
	static AES[] InvMixColumn(AES[] state) {
		byte[] b = { (byte) 0x0E, (byte) 0x09, (byte) 0x0D, (byte) 0x0B };
		AES a = new AES(b);
		for (int i = 0; i < Nb; i++)
			state[i] = multiply(a, state[i]);
		return state;
	}

	/**
	 * 轮密钥加变换
	 *
	 * @param state
	 * @param key
	 * @return
	 */
	static AES[] AddRoundKey(AES[] state, AES[] key) {
		for (int i = 0; i < Nb; i++)
			state[i] = add(state[i], key[i]);
		return state;
	}

	/**
	 * 加密
	 *
	 * @param plaintext 4个字长度的明文
	 * @param CipherKey 4、6或8个字长度的密钥
	 * @return
	 */
	public static AES[] encrypt(AES[] plaintext, AES[] CipherKey) {
		// 加密密钥扩展
		AES[][] RoundKey = KeyExpansion(CipherKey);
		AES[] ciphertext = new AES[plaintext.length];
		for (int i = 0; i < plaintext.length; i++)
			ciphertext[i] = new AES(plaintext[i]);
		// 初始轮密钥加
		ciphertext = AddRoundKey(ciphertext, RoundKey[0]);
		// 轮函数
		for (int i = 1; i < Nr + 1; i++) {
			// S盒变换
			ciphertext = ByteSub(ciphertext);
			// 行移位
			ciphertext = ShiftRow(ciphertext);
			// 列混合
			if (i != Nr)
				ciphertext = MixColumn(ciphertext);
			// 轮密钥加
			ciphertext = AddRoundKey(ciphertext, RoundKey[i]);
		}
		return ciphertext;
	}

	public static byte[] encryptFile(byte[] data, byte[] key) {
		return toByteArr(encrypt(toWordArr(data), toWordArr(key)));
	}

	public static byte[] encryptFile(byte[] data) {
		return toByteArr(encrypt(toWordArr(data), toWordArr(KEY.getBytes())));
	}

	public static void encryptFile(String sourceFile) throws Exception {
		File file = new File(sourceFile);
		long filesize = file.length();
		StringBuilder sb = new StringBuilder(sourceFile);
		sb.insert(sourceFile.lastIndexOf("."), "AES加密");
		String outFileName = sb.toString();
		System.out.println(outFileName);
		FileInputStream inputStream = new FileInputStream(sourceFile);
		FileOutputStream outputStream = new FileOutputStream(outFileName);
		double t1 = 0, t2 = 0, t3 = 0;
		byte[] bts = new byte[1048576];
		long dt1 = System.nanoTime();
		int readCount = inputStream.read(bts);
		long dt2 = System.nanoTime();
		while (readCount >= 0) {
			long dt3 = System.nanoTime();
			byte[] bx = encryptFile(bts, KEY.getBytes());
			long dt4 = System.nanoTime();
			outputStream.write(bx);
			long dt5 = System.nanoTime();
			readCount = inputStream.read(bts);
			long dt6 = System.nanoTime();

			t1 += dt6 - dt5;
			t2 += dt4 - dt3;
			t3 += dt5 - dt4;
		}
		t1 += dt2 - dt1;
		inputStream.close();
		outputStream.close();

		String datetime = new SimpleDateFormat("yyyy/MM/dd HH:mm:ss").format(new Date());
		appendToFile(datetime, 2, "Read", t1, sourceFile, filesize);
		appendToFile(datetime, 2, "Process", t2, sourceFile, filesize);
		appendToFile(datetime, 2, "Write", t3, sourceFile, filesize);
		appendToFile(datetime, 2, "Total", t1 + t2 + t3, sourceFile, filesize);
	}

	public static boolean appendToFile(String datetime, int type, String item, double time, String infile, long size)
			throws IOException {
		String info = String.format("%s,%d,%s,%.6f,%.6f,%d,%s\n", // 格式化为统一的CSV字符串
				datetime, type, item, // 时间、类型（0=加密，1=解密）和项目名称
				time / 1000000000.0, // 输入时间单位为ns，转换为s
				size * 1000000000.0 / 1073741824.0 / time, size, infile);
		String resfile = "results.csv";
		if (!new File(resfile).exists())
			appendToFile(resfile, "DateTime,Type,Command,Time(s),Speed(GB/s),Flesize(B),Filepath\n");
		appendToFile(resfile, info);
		return true;
	}

	public static void appendToFile(String filepath, String content) throws IOException {
		System.out.print(content);
		FileWriter writer = new FileWriter(filepath, true);
		writer.write(content);
		writer.close();
	}
	
	/**
	 * 解密
	 *
	 * @param ciphertext 4个字长度的密文
	 * @param CipherKey  4、6或8个字长度的密钥
	 * @return
	 */
	public static AES[] decrypt(AES[] ciphertext, AES[] CipherKey) {
		// 解密密钥扩展
		AES[][] InvRoundKey = InvKeyExpansion(CipherKey);
		AES[] plaintext = new AES[ciphertext.length];
		for (int i = 0; i < ciphertext.length; i++)
			plaintext[i] = new AES(ciphertext[i]);
		// 初始轮密钥加
		plaintext = AddRoundKey(plaintext, InvRoundKey[Nr]);
		// 轮函数
		for (int i = Nr - 1; i >= 0; i--) {
			// S盒变换
			plaintext = InvByteSub(plaintext);
			// 行移位
			plaintext = InvShiftRow(plaintext);
			// 列混合
			if (i != 0)
				plaintext = InvMixColumn(plaintext);
			// 轮密钥加
			plaintext = AddRoundKey(plaintext, InvRoundKey[i]);
		}
		return plaintext;
	}

	public static byte[] decryptFile(byte[] data, byte[] key) {
		return toByteArr(decrypt(toWordArr(data), toWordArr(key)));
	}

	public static byte[] decryptFile(byte[] data, String key) {
		return toByteArr(decrypt(toWordArr(data), toWordArr(key.getBytes())));
	}

	public static void decryptFile(String sourceFile) throws Exception {
		File file = new File(sourceFile);
		long filesize = file.length();
		String outFileName = sourceFile.replace("加密", "解密");
		FileInputStream inputStream = new FileInputStream(sourceFile);
		FileOutputStream outputStream = new FileOutputStream(outFileName);
		double t1 = 0, t2 = 0, t3 = 0;
		byte[] bts = new byte[1048576];
		long dt1 = System.nanoTime();
		int readCount = inputStream.read(bts);
		long dt2 = System.nanoTime();
		while (readCount != -1) {
			long dt3 = System.nanoTime();
			byte[] bx = decryptFile(bts, KEY.getBytes());
			long dt4 = System.nanoTime();
			outputStream.write(bx);
			long dt5 = System.nanoTime();
			readCount = inputStream.read(bts);
			long dt6 = System.nanoTime();

			t1 += dt6 - dt5;
			t2 += dt4 - dt3;
			t3 += dt5 - dt4;
		}
		t1 += dt2 - dt1;
		inputStream.close();
		outputStream.close(); 
		
		String datetime = new SimpleDateFormat("yyyy/MM/dd HH:mm:ss").format(new Date());
		appendToFile(datetime, 2, "Read", t1, sourceFile, filesize);
		appendToFile(datetime, 2, "Process", t2, sourceFile, filesize);
		appendToFile(datetime, 2, "Write", t3, sourceFile, filesize);
		appendToFile(datetime, 2, "Total", t1 + t2 + t3, sourceFile, filesize);
	}

	public static AES[] toWordArr(byte[] b) {
		int len = b.length / 4;
		if (b.length % 4 != 0)
			len++;
		AES[] w = new AES[len];
		for (int i = 0; i < len; i++) {
			byte[] c = new byte[4];
			if (i * 4 < b.length) {
				for (int j = 0; j < 4 && (i * 4 + j) < b.length; j++) {
					c[j] = b[i * 4 + j];
				}
			}
			w[i] = new AES(c);
		}
		return w;
	}

	public static byte[] toByteArr(AES[] w) {
		int len = w.length * 4;
		byte[] bytes = new byte[len];
		for (int i = 0; i < w.length; i++) {
			for (int j = 0; j < 4 && j < w[i].word.length; j++) {
				bytes[i * 4 + j] = w[i].word[j];
			}
		}
		return bytes;
	}

	public static String wordArrStr(AES[] w) {
		String str;
		StringBuilder sb = new StringBuilder();
		for (AES words : w) {
			str = words.toString();
			sb.append(str);
		}
		return sb.toString();
	} 

}
